

import tkinter as tk
from tkinter import filedialog
from tkinter import ttk
import shutil
import os
import hashlib
import mysql.connector
from cryptography.fernet import Fernet
import binascii
import base64

isUserChecked = False
files_list = []

# 用户登录函数
def checkUser():
    # 获取ID和密码
    user_id = user_id_entry.get()
    user_pwd = user_id_pw.get()

    # 查询记录
    query = "SELECT userid FROM user_info WHERE userid = %s AND userpwd = %s"
    cursor.execute(query, (user_id, user_pwd))
    result = cursor.fetchone()

    # 检查是否匹配
    global isUserChecked
    if result:
        isUserChecked = True
        status_label1.config(text="登录成功！")
    else:
        isUserChecked = False
        status_label1.config(text="登录失败！用户名或密码错误。")


def calculate_hash(file_path):
    # 计算文件哈希值
    hasher = hashlib.sha256()
    with open(file_path, 'rb') as f:
        while True:
            data = f.read(65536)  # 64 KB
            if not data:
                break
            hasher.update(data)
    return hasher.hexdigest()


def uploadFile():
    # 文件上传
    if not isUserChecked:
        status_label2.config(text="请先登录！")
        return

    # 获取用户输入的ID
    user_id = user_id_entry.get()

    file_path = filedialog.askopenfilename()
    if file_path:
        hash_value = calculate_hash(file_path)

        # 检查数据库是否存在相同哈希值的文件
        query = "SELECT ownerid, findex FROM file_info WHERE fhash = %s"
        cursor.execute(query, (hash_value,))
        existing_file = cursor.fetchone()
        cursor.fetchall()

        # 存在相同哈希值文件
        if existing_file:
            existing_owner_id, existing_file_path = existing_file
            # @
            if existing_owner_id == user_id:
                status_label2.config(text="文件已存在！")
                return
            else:
                insert_query = "INSERT INTO file_info (ownerid, findex, fhash) VALUES (%s, %s, %s)"
                cursor.execute(insert_query, (user_id, existing_file_path, hash_value))
                conn.commit()
        # 不存在相同哈希值文件
        else:
            # 如果不存在相同哈希值的文件，添加新记录
            destination_dir = os.path.join(r"E:\A课程资料\应用密码学\EP2", user_id)
            os.makedirs(destination_dir, exist_ok=True)

            key_bytes = binascii.unhexlify(hash_value)
            key_base64 = base64.urlsafe_b64encode(key_bytes)
            cipher_suite = Fernet(key_base64)

            with open(file_path, "rb") as file:
                data = file.read()
            encrypted_data = cipher_suite.encrypt(data)
            with open("aenc.txt", "wb") as file:
                file.write(encrypted_data)
            with open(os.path.join(destination_dir, os.path.basename(file_path)), "wb") as file:
                file.write(encrypted_data)

            # shutil.copy(file_path, destination_dir)
            final_file = os.path.join(destination_dir, os.path.basename(file_path))
            insert_query = "INSERT INTO file_info (ownerid, findex, fhash) VALUES (%s, %s, %s)"
            cursor.execute(insert_query, (user_id, final_file, hash_value))
            conn.commit()

        status_label2.config(text="文件上传成功")


def showFiles():
    global files_list

    if not isUserChecked:
        status_label2.config(text="请先登录！")
        return

    user_id = user_id_entry.get()

    query = "SELECT findex FROM file_info WHERE ownerid = %s"
    cursor.execute(query, (user_id,))
    files_list = [record[0] for record in cursor.fetchall()]

    # 显示到UI
    files_display.delete(0, tk.END)
    for file_path in files_list:
        file_name = os.path.basename(file_path)
        files_display.insert(tk.END, file_name)


def downloadFile(event):
    global files_list

    if not isUserChecked:
        status_label2.config(text="请先登录！")
        return

    # 获取选定的文件索引
    selected_index = files_display.curselection()
    if selected_index:
        selected_file_index = selected_index[0]
        file_path = files_list[selected_file_index]

        if not os.path.exists(file_path):
            status_label2.config(text="文件不存在：" + file_path)
            return

        default_file_name = os.path.basename(file_path)

        # 对话框选择文件保存路径
        save_path = filedialog.asksaveasfilename(
            initialdir="/",
            title="选择保存位置",
            filetypes=(("All files", "*.*"),),
            initialfile=default_file_name
        )

        if save_path:
            try:
                shutil.copy(file_path, save_path)
                print(file_path)
                print(save_path)

                # 从数据库中查询文件的哈希值
                query = "SELECT fhash FROM file_info WHERE findex = %s"
                cursor.execute(query, (file_path,))
                hash_result = cursor.fetchone()
                cursor.fetchall()

                print(hash_result)

                if hash_result:
                    hash_value = hash_result[0]
                    print(hash_value)

                    key_bytes = binascii.unhexlify(hash_value)
                    key_base64 = base64.urlsafe_b64encode(key_bytes)
                    cipher_suite = Fernet(key_base64)
                    with open(save_path, "rb") as encrypted_file:
                        encrypted_data = encrypted_file.read()
                    decrypted_data = cipher_suite.decrypt(encrypted_data)
                    print(decrypted_data)
                    with open(save_path, "wb") as decrypted_file:
                        decrypted_file.write(decrypted_data)

                status_label2.config(text="文件下载成功")
            except Exception as e:
                status_label2.config(text="文件下载失败：" + str(e))
        else:
            status_label2.config(text="未选择保存路径")
    else:
        status_label2.config(text="请选择要下载的文件")


# 连接到MySQL数据库
conn = mysql.connector.connect(
    host="localhost",
    user="root",
    password="tang1255184",
    database="Info"
)
# 创建游标对象
cursor = conn.cursor()
# UI界面
root = tk.Tk()
root.title("文件上传")
# 用户ID输入框
user_id_label = tk.Label(root, text="用户ID：")
user_id_label.grid(row=0, column=0, padx=5, pady=5)
user_id_entry = tk.Entry(root)
user_id_entry.grid(row=0, column=1, padx=5, pady=5)
# 用户密码输入框
user_id_pw = tk.Label(root, text="用户密码：")
user_id_pw.grid(row=1, column=0, padx=5, pady=5)
user_id_pw = tk.Entry(root)
user_id_pw.grid(row=1, column=1, padx=5, pady=5)
# 登录按钮
signin_button = tk.Button(root, text="登录", command=checkUser)
signin_button.grid(row=2, column=0, columnspan=1, padx=5, pady=5)
# 登陆状态
status_label1 = tk.Label(root, text="尚未登录！")
status_label1.grid(row=2, column=1, columnspan=2, padx=5, pady=5)
# 水平分割线
horizontal_separator = ttk.Separator(root, orient='horizontal')
horizontal_separator.grid(row=3, columnspan=2, sticky='ew', padx=10, pady=5)
# 上传按钮
upload_button = tk.Button(root, text="上传文件", command=uploadFile)
upload_button.grid(row=4, column=0, padx=5, pady=5)
# 查看按钮
upload_button = tk.Button(root, text="查看文件", command=showFiles)
upload_button.grid(row=4, column=1, padx=5, pady=5)
# 文件显示列表
files_display = tk.Listbox(root,width=30)
files_display.grid(row=5, column=0, columnspan=2, padx=5, pady=5)
# 状态提示
status_label2 = tk.Label(root, text="当前尚未执行文件操作")
status_label2.grid(row=6, column=0, columnspan=2, padx=5, pady=5)
# 绑定双击事件处理函数
files_display.bind("<Double-Button-1>", downloadFile)
root.mainloop()
# 关闭游标和连
cursor.close()
conn.close()


